// File provided for Microsoft Shared Source.
// Copyright (c) Microsoft Corporation. All rights reserved.

#include "StdAfx.h"
#include "WMIEnumerator.h"
#include "CWMIFldr.h"

HRESULT CWMIEnumerator::_Initialize(__in CWMIEnumerator *pWMIEnum)
{
    HRESULT hr;
    hr = pWMIEnum->_spEnum->Clone(&_spEnum);
    TRACEFNEXIT();
    return hr;
}

HRESULT CWMIEnumerator::_Initialize(__in WMIENUMINFO *pEnumInfo)
{
    HRESULT hr;

    CComPtr<CMultiEnumerator> spMultiEnum;
    hr = CMultiEnumerator::Create(&spMultiEnum);
    if (SUCCEEDED(hr))
    {
        CComPtr<IWbemServices> spSvc;
        hr = pEnumInfo->GetServices(&spSvc);
        if (SUCCEEDED(hr))
        {
            CComBSTR sbstrLanguage;
            hr = sbstrLanguage.Append(L"WQL");
            if (SUCCEEDED(hr))
            {
                CComBSTR sbstrClassName;
                hr = sbstrClassName.Append(pEnumInfo->pszClass);
                if (SUCCEEDED(hr))
                {
                    if (!pEnumInfo->pszClass)
                    {
                        CComBSTR sbstrNSQuery;
                        hr = sbstrNSQuery.Append(L"SELECT * FROM __NAMESPACE");
                        if (SUCCEEDED(hr))
                        {
                            CComPtr<IEnumWbemClassObject> spEnumwcoNamespaces;
                            hr = spSvc->ExecQuery(sbstrLanguage, sbstrNSQuery,
                                WBEM_FLAG_BIDIRECTIONAL, // don't specify forward-only
                                0, // Context
                                &spEnumwcoNamespaces);
                            if (SUCCEEDED(hr))
                            {
                                hr = spMultiEnum->AddEnumerator(spEnumwcoNamespaces);
                            }
                        }
                    }
                    else
                    {
                        CComPtr<IEnumWbemClassObject> spEnumwcoInstances;
                        hr = spSvc->CreateInstanceEnum(sbstrClassName,
                            WBEM_FLAG_BIDIRECTIONAL | WBEM_FLAG_SHALLOW, // don't specify forward-only
                            0, // Context
                            &spEnumwcoInstances);
                        if (SUCCEEDED(hr))
                        {
                            hr = spMultiEnum->AddEnumerator(spEnumwcoInstances);
                        }
                    }
                }
                if (SUCCEEDED(hr))
                {
                    CComPtr<IEnumWbemClassObject> spEnumwcoClasses;
                    hr = spSvc->CreateClassEnum(sbstrClassName,
                        WBEM_FLAG_BIDIRECTIONAL, // don't specify forward-only
                        0, // Context
                        &spEnumwcoClasses);
                    if (SUCCEEDED(hr))
                    {
                        hr = spMultiEnum->AddEnumerator(spEnumwcoClasses);
                    }
                }
            }
        }
        if (SUCCEEDED(hr))
        {
            _spEnum.Attach(spMultiEnum.Detach());
        }
    }

    TRACEFNEXIT();
    return hr;
}

CWMIEnumerator::CWMIEnumerator(void)
{
    InterlockedIncrement(&g_cDllLock);

    _cRef = 1;
}

CWMIEnumerator::~CWMIEnumerator(void)
{
    InterlockedDecrement(&g_cDllLock);
}

HRESULT CWMIEnumerator::QueryInterface(REFIID riid, void **ppv)
{
    //TRACEFNENTER_IID(riid);

    HRESULT hr = E_NOINTERFACE;
    *ppv = NULL;

    QI_ENTRY(IUnknown);
    QI_ENTRY(IEnumIDList);

    TRACEFNEXIT_IID(riid);
    return hr;
}

HRESULT CWMIEnumerator::Next(ULONG celt, PITEMID_CHILD *rgelt, ULONG *pceltFetched)
{
    TRACEFNENTER();

    HRESULT hr = S_OK;

    // Retrieve the objects in the result set.
    DWORD iPidl = 0;
    CAtlArray< CComPtr<IWbemClassObject> > rgspWMIObj;
    while (SUCCEEDED(hr) && iPidl == 0)
    {
        rgspWMIObj.SetCount(0);
        rgspWMIObj.SetCount(celt);

        DWORD cel = 0;
        hr = _spEnum->Next(
            0,                  // Time out
            celt,               // object count requested
            &rgspWMIObj[0], &cel);

        if (SUCCEEDED(hr))
        {
            if (cel > 0)
            {
                rgspWMIObj.SetCount(cel);
                for (DWORD iObj = 0; iObj < cel; iObj++)
                {
                    hr = CWMIFldr::s_CreateIDList(rgspWMIObj[iObj], &rgelt[iPidl]);
                    if (SUCCEEDED(hr))
                    {
                        iPidl++;
                    }
                    else
                    {
                        cel--;
                        hr = S_OK;
                    }
                }
            }
            
            if (S_FALSE == hr)
            {
                break;
            }
        }
    }

    if (pceltFetched)
    {
        *pceltFetched = iPidl;
    }

    if (FAILED(hr))
    {
        hr = E_FAIL;
    }

    TRACEFNEXIT();
    return hr;
}

HRESULT CWMIEnumerator::Reset()
{
    TRACEFNENTER();

    HRESULT hr = _spEnum->Reset();

    TRACEFNEXIT();
    return hr;
}

HRESULT CWMIEnumerator::Clone(IEnumIDList **ppEnum)
{
    TRACEFNENTER();
    HRESULT hr = CWMIEnumerator::Create(this, IID_PPV_ARGS(ppEnum));
    TRACEFNEXIT();
    return hr;
}

HRESULT CWMIEnumerator::Skip(ULONG)
{
    TRACEFNENTER();

    HRESULT hr = E_NOTIMPL;

    TRACEFNEXIT();
    return hr;
}

////////////////////////
// CMultiEnumerator
//////////

CMultiEnumerator::CMultiEnumerator()
{
    InterlockedIncrement(&g_cDllLock);

    _cRef = 1;
    _iEnumCurrent = 0;
}

CMultiEnumerator::~CMultiEnumerator()
{
    InterlockedDecrement(&g_cDllLock);
}

HRESULT CMultiEnumerator::QueryInterface(REFIID riid, __out void **ppv)
{
    HRESULT hr = E_NOINTERFACE;
    *ppv = NULL;

    QI_ENTRY_EX(IUnknown, IEnumWbemClassObject);
    QI_ENTRY(IEnumWbemClassObject);

    TRACEFNEXIT_IID(riid);
    return hr;
}

HRESULT CMultiEnumerator::Create(__deref_out CMultiEnumerator **ppMultiEnum)
{
    *ppMultiEnum = new CMultiEnumerator();
    return *ppMultiEnum ? S_OK : E_OUTOFMEMORY;
}

HRESULT CMultiEnumerator::AddEnumerator(__in IEnumWbemClassObject *pEnumWbemClassObject)
{
    // delete this line
    HRESULT hr = S_OK;
    try 
    { 
        _srgpEnumerators.Add(pEnumWbemClassObject);
    } 
    catch (CAtlException &exception) 
    {
        hr = exception; 
    } 
    return hr;
}

// IEnumWbemClassObject
HRESULT CMultiEnumerator::Reset()
{
    HRESULT hr = E_INVALIDARG;
    while (true)
    {
        hr = _srgpEnumerators[_iEnumCurrent]->Reset();
        if (_iEnumCurrent != 0)
        {
            _iEnumCurrent--;
        }
        else
        {
            break;
        }
    }
    return hr;
}

HRESULT CMultiEnumerator::Next(__in long lTimeout, __in ULONG uCount, __out_ecount_part(uCount, *puReturned) IWbemClassObject **apObjects, __out ULONG *puReturned)
{
    HRESULT hr = S_FALSE;
    
    while (true)
    {
        if (_iEnumCurrent < _srgpEnumerators.GetCount())
        {
            hr = _srgpEnumerators[_iEnumCurrent]->Next(lTimeout, uCount, apObjects, puReturned);
            if (S_FALSE == hr)
            {
                _iEnumCurrent++;
            }
            else
            {
                break;
            }
        }
        else
        {
            break;
        }
    }

    TRACEFNEXIT();
    return hr;
}

HRESULT CMultiEnumerator::NextAsync(__in ULONG uCount, __in_opt IWbemObjectSink *pSink)
{
    ATLASSERT(FALSE);
    return E_NOTIMPL;
}

HRESULT CMultiEnumerator::Clone(__deref_out IEnumWbemClassObject **ppEnum)
{
    CAtlArray< CComPtr<IEnumWbemClassObject> > srgpEnum;
    HRESULT hr = srgpEnum.SetCount(_srgpEnumerators.GetCount()) ? S_OK : E_OUTOFMEMORY;
    if (SUCCEEDED(hr))
    {
        for (size_t iEnum = 0; SUCCEEDED(hr) && iEnum < _srgpEnumerators.GetCount(); iEnum++)
        {
            hr = _srgpEnumerators[iEnum]->Clone(&srgpEnum[iEnum]);
        }

        if (SUCCEEDED(hr))
        {
            CComPtr<CMultiEnumerator> spMultiEnumerator;
            hr = CMultiEnumerator::Create(&spMultiEnumerator);
            for (size_t iEnum = 0; SUCCEEDED(hr) && iEnum < srgpEnum.GetCount(); iEnum++)
            {
                hr = spMultiEnumerator->AddEnumerator(srgpEnum[iEnum]);
            }
            if (SUCCEEDED(hr))
            {
                *ppEnum = spMultiEnumerator.Detach();
            }
        }
    }
    return hr;
}

HRESULT CMultiEnumerator::Skip(__in long lTimeout, __in ULONG nCount)
{
    ATLASSERT(FALSE);
    return E_NOTIMPL;
}

///////////////////////////
// WMIENUMINFO
/////////

WMIENUMINFO::WMIENUMINFO()
{
    pszClass = NULL;
    pszComputer = NULL;
    pszInstance = NULL;
    pszNamespace = NULL;
}

WMIENUMINFO::~WMIENUMINFO()
{
}

HRESULT WMIENUMINFO::GetServices(IWbemServices **ppServices)
{
    return CWMIFldr::s_GetServer(pszComputer, pszNamespace, ppServices);
}

HRESULT WMIENUMINFO::GetInstancePath(BOOL fAbsolute, BSTR *pbstrInstancePath)
{
    HRESULT hr = S_OK;
    CComBSTR sbstrInstancePath;
    if (fAbsolute)
    {
        hr = sbstrInstancePath.Append(L"\\\\");
        if (SUCCEEDED(hr))
        {
            hr = sbstrInstancePath.Append(pszComputer);
            if (SUCCEEDED(hr))
            {
                hr = sbstrInstancePath.Append(L"\\");
                if (SUCCEEDED(hr))
                {
                    hr = sbstrInstancePath.Append(pszNamespace);
                    if (SUCCEEDED(hr))
                    {
                        hr = sbstrInstancePath.Append(L":");
                    }
                }
            }
        }
    }

    if (SUCCEEDED(hr))
    {
        hr = sbstrInstancePath.Append(pszClass);
        if (SUCCEEDED(hr))
        {
            hr = sbstrInstancePath.Append(".");
            if (SUCCEEDED(hr))
            {
                hr = sbstrInstancePath.Append(pszInstance);
            }
        }
    }

    if (SUCCEEDED(hr))
    {
        *pbstrInstancePath = sbstrInstancePath.Detach();
    }
    else
    {
        *pbstrInstancePath = NULL;
    }

    return hr;
}

HRESULT WMIENUMINFO::GetInstance(IWbemClassObject **ppInstance)
{
    CComPtr<IWbemServices> spSvc;
    HRESULT hr = GetServices(&spSvc);
    if (SUCCEEDED(hr))
    {
        hr = GetInstance(spSvc, ppInstance);
    }
    return hr;
}

HRESULT WMIENUMINFO::GetInstance(IWbemServices *pSvc, IWbemClassObject **ppInstance)
{
    CComBSTR sbstrInstancePath;
    HRESULT hr = GetInstancePath(FALSE, &sbstrInstancePath);
    if (SUCCEEDED(hr))
    {
        hr = pSvc->GetObjectW(sbstrInstancePath,  WBEM_FLAG_RETURN_WBEM_COMPLETE, NULL, ppInstance, NULL);
    }
    return hr;
}

HRESULT WMIENUMINFO::InitFromPath(__in PCWSTR pszParsingPath)
{
    HRESULT hr;

    CComPtr<IWbemPath> spPathParser;
    hr = CoCreateInstance(__uuidof(WbemDefPath), 0, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&spPathParser));
    if (SUCCEEDED(hr))
    {
        hr = spPathParser->SetText(WBEMPATH_CREATE_ACCEPT_ALL, pszParsingPath);
        if (SUCCEEDED(hr) || WBEM_E_INVALID_PARAMETER == hr)
        {
            ULONG cchServer = ARRAYSIZE(szServer);
            hr = spPathParser->GetServer(&cchServer, szServer);
            if (SUCCEEDED(hr))
            {
                pszComputer = szServer;
            }
            else if (WBEM_E_NOT_AVAILABLE == hr)
            {
                hr = S_OK;
            }

            if (SUCCEEDED(hr))
            {
                ULONG cNamespaces;
                hr = spPathParser->GetNamespaceCount(&cNamespaces);
                if (SUCCEEDED(hr))
                {
                    ULONG cchNamespaceRemaining = ARRAYSIZE(szNamespace);
                    PWSTR szNamespaceRemaining = szNamespace;
                    for (ULONG iNamespace = 0; SUCCEEDED(hr) && iNamespace < cNamespaces; iNamespace++)
                    {
                        if (iNamespace != 0)
                        {
                            hr = StringCchCopy(szNamespaceRemaining, cchNamespaceRemaining, L"\\");
                            cchNamespaceRemaining--;
                            szNamespaceRemaining++;
                        }

                        if (SUCCEEDED(hr))
                        {
                            ULONG cchNamespace = cchNamespaceRemaining;
                            hr = spPathParser->GetNamespaceAt(iNamespace, &cchNamespace, szNamespaceRemaining);
                            if (SUCCEEDED(hr))
                            {
                                cchNamespace--; // they count the null, we don't want to
                                cchNamespaceRemaining -= cchNamespace;
                                szNamespaceRemaining += cchNamespace;
                            }
                        }
                    }
                }
                if (SUCCEEDED(hr))
                {
                    pszNamespace = szNamespace;
                }
                else if (WBEM_E_NOT_AVAILABLE == hr)
                {
                    hr = S_OK;
                }

                if (SUCCEEDED(hr))
                {
                    ULONG cchClassName = ARRAYSIZE(szClassName);
                    hr = spPathParser->GetClassName(&cchClassName, szClassName);
                    if (SUCCEEDED(hr))
                    {
                        pszClass = szClassName;
                    }
                    else if (WBEM_E_NOT_AVAILABLE == hr)
                    {
                        hr = S_OK;
                    }
                }

                BOOL fIsNamespace = pszClass 
                    && CompareString(LOCALE_INVARIANT, NORM_IGNORECASE, pszClass, -1, L"__NAMESPACE", -1) == CSTR_EQUAL;

                if (fIsNamespace)
                {
                    // this is a namespace and we want to treat it as a folder, so lets yank the class name and append the instance name to the namespace
                    pszClass = NULL;
                }

                if (SUCCEEDED(hr))
                {
                    CComPtr<IWbemPathKeyList> spKeyList;
                    hr = spPathParser->GetKeyList(&spKeyList);
                    if (SUCCEEDED(hr))
                    {
                        if (!fIsNamespace)
                        {
                            ULONG cchInstance = ARRAYSIZE(szInstance);
                            hr = spKeyList->GetText(WBEMPATH_QUOTEDTEXT, &cchInstance, szInstance);
                            if (SUCCEEDED(hr) && wcslen(szInstance) > 0)
                            {
                                pszInstance = szInstance;
                            }
                        }
                        else
                        {
                            size_t cchRemaining;
                            PWSTR pszNamespaceRemaining;
                            hr = StringCchCatEx(szNamespace, ARRAYSIZE(szNamespace), L"\\", &pszNamespaceRemaining, &cchRemaining, 0);
                            if (SUCCEEDED(hr))
                            {
                                ULONG cchKeyName = 0;
                                CComVariant svarName;
                                CIMTYPE cimType;
                                hr = spKeyList->GetKey2(0, 0, &cchKeyName, NULL, &svarName, (ULONG *)&cimType);
                                if (SUCCEEDED(hr))
                                {
                                    hr = VariantToString(svarName, pszNamespaceRemaining, (UINT)cchRemaining);
                                }
                            }
                        }
                    }
                    else if (WBEM_E_INVALID_PARAMETER == hr)
                    {
                        // not sure why they didn't stick to WBEM_E_NOT_AVAILABLE
                        hr = S_OK;
                    }
                }
            }
        }
    }

    return hr;
}

//////////////////////////////
// CWMIInstanceEnumerator
///////////

HRESULT CWMIInstanceEnumerator::_Initialize(WMIENUMINFO *pEnumInfo)
{
    HRESULT hr = pEnumInfo->GetInstance(&_spInstance);
    if (SUCCEEDED(hr))
    {
        // Get the property names from this instance.
        SAFEARRAY *psaNames = NULL;
        hr = _spInstance->GetNames(NULL, WBEM_FLAG_ALWAYS, NULL, &psaNames);
        if (SUCCEEDED(hr))
        {
            // Get the number of properties.
            long lLower, lUpper; 
            hr = SafeArrayGetLBound(psaNames, 1, &lLower);
            if (SUCCEEDED(hr))
            {
                hr = SafeArrayGetUBound(psaNames, 1, &lUpper);
                if (SUCCEEDED(hr))
                {
                    hr = _asbstrPropertyNames.SetCount(lUpper - lLower + 1) ? S_OK : E_OUTOFMEMORY;
                    if (SUCCEEDED(hr))
                    {
                        for (long i = lLower; SUCCEEDED(hr) && i <= lUpper; i++) 
                        {
                            BSTR bstrPropNameWeak;
                            hr = SafeArrayGetElement(psaNames, &i, &bstrPropNameWeak);
                            if (SUCCEEDED(hr))
                            {
                                hr = _asbstrPropertyNames[i-lLower].AppendBSTR(bstrPropNameWeak);
                            }
                        }
                    }
                }
            }
            SafeArrayDestroy(psaNames);
        }
    }
    return hr;
}

HRESULT CWMIInstanceEnumerator::QueryInterface(REFIID riid, void **ppv)
{
    //TRACEFNENTER_IID(riid);

    HRESULT hr = E_NOINTERFACE;
    *ppv = NULL;

    QI_ENTRY(IUnknown);
    QI_ENTRY(IEnumIDList);

    TRACEFNEXIT_IID(riid);
    return hr;
}

HRESULT CWMIInstanceEnumerator::Next(ULONG celt, PITEMID_CHILD *rgelt, ULONG *pceltFetched)
{
    TRACEFNENTER();

    HRESULT hr = S_OK;
    DWORD iPidl = 0;
    for (; iPidl < celt && _iProperty < _asbstrPropertyNames.GetCount(); _iProperty++)
    {
        if (SUCCEEDED(CWMIFldr::s_CreateIDList(_spInstance, _asbstrPropertyNames[_iProperty], &rgelt[iPidl])))
        {
            iPidl++;
        }
    }

    if (0 == iPidl)
    {
        hr = S_FALSE;
    }

    if (pceltFetched)
    {
        *pceltFetched = iPidl;
    }

    TRACEFNEXIT();
    return hr;
}

HRESULT CWMIInstanceEnumerator::Reset()
{
    TRACEFNENTER();
    HRESULT hr = S_OK;

    _iProperty = 0;

    TRACEFNEXIT();
    return hr;
}

HRESULT CWMIInstanceEnumerator::Clone(IEnumIDList **ppEnum)
{
    TRACEFNENTER();

    HRESULT hr = E_NOTIMPL;

    TRACEFNEXIT();
    return hr;
}

HRESULT CWMIInstanceEnumerator::Skip(ULONG)
{
    TRACEFNENTER();

    HRESULT hr = E_NOTIMPL;

    TRACEFNEXIT();
    return hr;
}
